import json
import random

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from datetime import datetime

def plot_gantt(json_file):
    # 加载JSON数据
    with open(json_file, 'r') as f:
        data = json.load(f)

    moves = data["MoveList"]


    # 生成随机颜色映射，用于工件ID
    def generate_color():
        return f'#{random.randint(0, 0xFFFFFF):06x}'  # 生成一个随机的颜色

    # 数据预处理
    modules = sorted(list(set([m["ModuleName"] for m in moves])))    # 机器列表
    module_ypos = {module: i for i, module in enumerate(modules)}

    # 为每个工件生成一个颜色映射
    matid_colors = {}
    for move in moves:
        matid = move["MatID"]
        if matid not in matid_colors:
            matid_colors[matid] = generate_color()

    # 创建画布
    plt.figure(figsize=(20, 10))

    # 设置横轴范围
    max_time = max([move["EndTime"] for move in moves])
    min_time = min([move["StartTime"] for move in moves])

    # 绘制每个工序
    for move in moves:
        module = move["ModuleName"]
        duration = move["EndTime"] - move["StartTime"]
        matid = move["MatID"]

        # 为工件ID选择颜色
        color = matid_colors.get(matid)  # 默认颜色是灰色

        # 绘制条形
        plt.barh(module_ypos[module],
                 duration,
                 left=move["StartTime"],
                 height=0.3,
                 color=color,
                )

    # 坐标轴设置
    plt.yticks(range(len(modules)), modules)    # y轴刻度位置
    plt.xlabel('Time (seconds)')
    plt.ylabel('Modules')
    plt.title('Scheduling Gantt Chart-5a-final-test-improve-new')


    # 添加图例
    legend_patches = [
        mpatches.Patch(color=color, label=f'MatID {matid}')
        for matid, color in matid_colors.items()
    ]

    # 将图例添加到图表中
    plt.legend(handles=legend_patches, bbox_to_anchor=(0.5, -0.15), loc='upper center', ncol=10)


    # 调整布局
    plt.grid(axis='x', alpha=0.5)
    plt.tight_layout(pad=1.0)  # 调整间距

    # 保存图像
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    plt.savefig(f'gantt_{timestamp}.png', dpi=300, bbox_inches='tight')
    plt.close()  # 关闭图形，避免内存占用


# 使用示例
if __name__ == "__main__":
    plot_gantt(r'D:\CodeProject\PyCode\ChipSchduile\res\5a-final-test-improve(2).json')